import os
from pathlib import Path
import numpy as np


# 数据相关
# !!!!!!!!!!!!!!!!!!!!!!!!
# source_root = Path(r'/public/pazhou/pazhou_data/train')
source_root = Path(r'D:\jassorRepository\emergency')
image_root = source_root / 'data'
label_seg_root = source_root / 'mask'
label_cls_path = source_root / 'train.csv'

images = {fname.removesuffix('.nii.gz'): image_root / f'' for fname in os.listdir(image_root)}
labels_seg = {fname.removesuffix('.nii.gz'): label_seg_root / f'' for fname in os.listdir(label_seg_root)}
with open(label_cls_path, 'r') as f:
    f.readline()
    labels_cls = [line.split(',') for line in f.readlines()]
    labels_cls = {name: tuple(map(int, [liver, spleen, left_kidney, right_kidney])) for name, liver, spleen, left_kidney, right_kidney in labels_cls}

# 数据集相关
available = list(images.keys())
trains = list(np.random.choice(available, 300))
valids = list(set(available) - set(trains))

# 训练相关
# !!!!!!!!!!!!!!!!!!!!!!!!
device = 'cuda:0'
# batch_size = 16
batch_size = 2
# train_epoch = 100
train_epoch = 1

# 输出相关
workspace = Path('..') if os.path.basename(os.path.abspath('.')) in ['dev', 'test', 'main'] else Path('.')
target = 'latest'
weight_root = workspace / 'output' / target / 'weights'
log_path = workspace / 'output' / target / 'log.txt'
tensorboard_root = workspace / 'output' / target / 'tensorboard'
metrics_root = workspace / 'output' / target / 'metric'
key_path = workspace / 'output' / target / 'key_set.txt'
config_backup_path = workspace / 'output' / target / 'config.py'
os.makedirs(workspace / 'output' / target, exist_ok=True)
os.makedirs(metrics_root, exist_ok=True)
metric_csv_sep = ','
